module (...,package.seeall) 

gr = require "graphs"
pq = require "priority_queues"


-- set up an "IDT" table associating each node 
-- to its *in-degree*, i.e. the number of 
-- edges pointing into it.  

function idt_of_graph(g)
   local idt = {} 
   -- initialize:  No edges yet.
   for i = 1,g.length do idt[i]=0 end 

   for i = 1,g.length do 
      for _,e in ipairs(g[i]) do 
	 -- count this new edge
	 -- into destination
	 idt[e.dst] = idt[e.dst] + 1
      end
   end 
   return idt 
end 

-- what to do to the IDT when removing a node 
-- from further consideration 

function update_idt(g,idt,removed_node,fix) 
   for _,e in ipairs(g[removed_node]) do 
      local j = e.dst 
      -- count down if a node that contributed an edge to j 
      -- is removed from the priority queue
      idt[j] = idt[j] - 1
      fix(j)			-- do further fixups, e.g. heapify up j 
   end 
end 

-- A location table is 
-- intended to map an item to its 
-- priority queue location 
-- (if the latter non-nil).  

-- Given an array, return a function 
-- that will install an updated location 
-- for each item.  

function loc_tab_make_loc_fn(loc_tab) 
   local function loc(v,new_l) 
      if v then loc_tab[v] = new_l end 
   end
   return loc 
end 

-- Given an IDT, return a fn that will 
-- serve as a priority queue preference 
-- function.  In particular, a value is 
-- preferred if it has a larger IDT value, 
-- or if its IDT value is non-nil and the 
-- other is nil.  

comparisonCount = 0

function prefer_by_in_degree(idt) 
   local function pref(a,b) 
      local ia = idt[a] 
      local ib = idt[b] 
	  
	  comparisonCount = comparisonCount + 1
	  
      if ia and ib 
      then return ia<ib 
      else return ia 
      end 
   end 
   return pref
end 

-- Two debugging functions:  Print contents of an idt, 
-- and of the heap.  

function print_idt(g,ts)
   count = 1
   for i = 1,g.length do 
      if count == 4 then io.write("\n") count = 1 else count = count+1 end 
      io.write(string.format("Node %d indeg: %d     ", i, ts.idt[i]))
   end 
   io.write("\n")
end 

function print_heap(g,ts) 
   local width = 1 
   local this_line = 0 
   for i = 1,g.length do 
      if this_line == width 
      then io.write("\n") ; width = 2 * width ; this_line = 1
      else this_line = this_line+1
      end 
      io.write(string.format("%d indeg: %d     ", ts.hp[i] or 1000, 
			     ts.idt[ts.hp[i]] or 1000))
   end 
   io.write("\n")
   io.flush() 
end 

-- A topological sorting table ts_tab for graph g 
-- is a triple with
--     a heap used as a priority queue, 
--     an IDT for g, and 
--     a location table.  
-- The heap is arranged to use a preference function 
-- that consults the IDT, and it is initialized to 
-- contain all the nodes of g.


function make_ts_tab_for(g) 
   local idtab = idt_of_graph(g)
   local loc_tab = {} 
   local locate = loc_tab_make_loc_fn(loc_tab)
   local h = pq.build_heap({},
			   prefer_by_in_degree(idtab),
			   locate)
   for i = 1,g.length do 
      pq.insert(h,i)
   end 

   return { hp = h, idt = idtab, lt = loc_tab }
end

-- If ts_tab is a topological sorting table for g, 
-- and a is an array to store the successive 
-- values in, then this fn removes the top (minimal) 
-- element min from the heap ts_tab.hp.  
-- It inserts it in the next position in array a.
-- It must also update the IDT to reflect removal of min. 

-- Trickiest point:  When decreasing the IDT entry for 
-- any m that min has an edge to, we must heapify up for it; 
-- since it is at heap location ts_tab.lt[m], that is the  
-- heapify-up second argument.  

-- Non-zero minimum implies some cyclicity.  

function get_minimal(g,ts_tab,a) 
   local min = pq.extract_top(ts_tab.hp) 
   if not(min) 
   then 
      return false 
   else if ts_tab.idt[min] == 0 
      then 
	 table.insert(a,min) 
	 update_idt(g,ts_tab.idt,min,
		    function (m) pq.heapify_up(ts_tab.hp,ts_tab.lt[m]) end)
	 return true
      else if type(min) == "number" 
	 then 
	    error(string.format("Non-zero indegree to %d:  Cycle exists.", min))
	 else
	    print(min) 
	    error("Weird value in table")
	 end 
      end 
   end 
end 

-- Now all we need to do is to make a ts_tab and 
-- get_minimal till there is nothing more to get.  

function topological_sort (g) 
   local ts = make_ts_tab_for(g)
   local a = {} 

   while (true == get_minimal(g,ts,a)) do end
   return a 
end 

-- find the position in array at which v resides 
function index_of_value(array,v) 
   local result = 0 
   for i = 1,#array do 
      if array[i] == v 
      then 
	 result = i
      end
   end 
   if result == 0 then error(string.format("Value %d not found.", v)) end 
   return result 
end 

-- test whether array is a topological sort of g.  
-- Check that all the arrows in g point from a 
-- node at an earlier index in array to a node 
-- at a later index.  

function test_topological_sort(g,array) 
   local function up_array(a,b) 
      return index_of_value(array,a)<index_of_value(array,b)
   end 

   for i = 1,g.length do 
      for _,e in ipairs(g[i]) do 
	 if not(up_array(i,e.dst)) 
	 then 
	    error(string.format("Edge from %d to %d isn't up.",i,e.dst))
	 end 
      end
   end
   return true
end 

function sort_n_test(g)
   local a = topological_sort (g)
   test_topological_sort(g,a)
   print(".")
end 
